#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
    根据因子生成器里的因子文件批量生产因子
'''
from typing import Tuple
import pdb
import os
import sys
from joblib import Parallel, delayed
from datetime import datetime
import pandas as pd
import numpy as np

_ = os.path.abspath(os.path.dirname(__file__))  # 返回当前文件路径
root_path = os.path.abspath(os.path.join(_, '..'))  # 返回根目录文件夹

sys.path.append(root_path)

from alphalib.config_backtest import output_path
from alphalib.contrib.strategy import z1_conf_types, z1_strategy, z1_playback

def get_strategy(strategy_name):
    stg_list = []
    def weight_func() -> Tuple[list, np.ndarray]:
        bh_list = [120, 360, 1080,]
        return bh_list, np.ones(len(bh_list))

    for n, w in zip(*weight_func()):
        strategy_conf = z1_strategy.Z1StrategyConfig(
            strategy_name=f'{strategy_name}_bh_{n}',
            hold_period = 1,
            long_factors=[z1_conf_types.F1FactorConfig(strategy_name, False, n, 0, 1)],
            short_factors=[z1_conf_types.F1FactorConfig(strategy_name, False, n, 0, 1)],
            filter_before_params=[],
            filter_after_params=[],
            if_use_spot=True,
            long_weight=2,
            short_weight=0,
            long_coin_num=.15,
            short_coin_num=0,
        )
        stg_list.append(z1_strategy.Z1Strategy(strategy_conf))

    strategy = z1_strategy.Z1MultiStrategy(stg_list)
    return strategy


def training_backtest(strategy_name):
    stg = get_strategy(strategy_name)
    z1_conf = z1_playback.Z1PlaybackConfig(
        compound_name=strategy_name,
        start_date='2020-01-01',
        end_date='2024-07-01',
        leverage=1,          # 杠杆建议定死在 1
        enable_funding_rate=False
    )

    res, curve, account_df, order_df = z1_playback.run_playback(stg, z1_conf)

    save_dir = os.path.join(output_path, strategy_name)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    equity_res_path = os.path.join(save_dir, '净值持仓数据.csv')
    res.to_csv(equity_res_path, encoding='gbk')
    curve.to_csv(equity_res_path, encoding='gbk', mode='a')
    # account_df.to_csv(os.path.join(save_dir, '虚拟账户数据.csv'), encoding='gbk')
    # order_df.to_pickle(os.path.join(save_dir, '下单面板数据.pkl'))

    sumarize_path = os.path.join(output_path, '训练结果统计.csv')
    if not os.path.exists(sumarize_path):
        res.to_csv(sumarize_path, encoding='gbk')
    else:
        res.to_csv(sumarize_path, encoding='gbk', mode='a', header=False)

if __name__ == '__main__':
    
    factor_df = pd.read_csv(os.path.join(root_path, './output/factors.csv'))
    batch_size = 10
    batch = []
    for index, factor_item in factor_df.iterrows():
        if index < 0:
            continue
        print(f'======================= index: {index} ========================== \n')
        strategy_name = factor_item['factor']
        batch.append(strategy_name)

        if len(batch) == batch_size:
            Parallel(n_jobs=batch_size)(
                delayed(training_backtest)(strategy) for strategy in batch
            )
            batch = []

        # training_backtest(strategy_name)
    # 处理剩余的任务
    if batch:
        Parallel(n_jobs=-1)(
            delayed(training_backtest)(strategy) for strategy in batch
        )

